import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import os
from matplotlib.ticker import MaxNLocator

# Set matplotlib to non-interactive backend to avoid Tkinter errors
plt.switch_backend('Agg')

# Configure font settings for proper Chinese character display
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

class ModelComparisonVisualizer:
    """Model training comparison visualizer for mix and single loss methods"""
    
    def __init__(self, mix_csv_path, single_csv_path, output_dir='outputs/comparison', max_epoch=None):
        self.mix_csv_path = mix_csv_path
        self.single_csv_path = single_csv_path
        self.output_dir = output_dir
        
        # Create output directory if it doesn't exist
        os.makedirs(self.output_dir, exist_ok=True)
        
        # Load and preprocess data
        self.mix_data = self._load_and_preprocess_data(mix_csv_path)
        self.single_data = self._load_and_preprocess_data(single_csv_path)
        
        # 自动检测共同的metrics和epoch范围
        self.common_metrics = self._detect_common_metrics()
        self.common_min_epoch = self._detect_common_min_epoch()
        self.common_max_epoch = self._detect_common_max_epoch()
        
        # 如果用户没有指定max_epoch，使用共同的最大epoch
        self.max_epoch = max_epoch if max_epoch is not None else self.common_max_epoch
        
        # 进一步筛选数据到共同的epoch范围
        self._filter_data_by_common_epochs()
        
        print(f"自动检测结果：")
        print(f"  共同最小epoch: {self.common_min_epoch}")
        print(f"  共同最大epoch: {self.common_max_epoch}")
        print(f"  使用的最大epoch: {self.max_epoch}")
        print(f"  检测到的共同metrics数量: {len(self.common_metrics)}")
    
    def _load_and_preprocess_data(self, csv_path):
        """Load and preprocess data from CSV file"""
        try:
            # Read CSV file
            df = pd.read_csv(csv_path)
            return df
        except Exception as e:
            print(f"Error loading data from {csv_path}: {e}")
            raise
    
    def _detect_common_metrics(self):
        """自动检测两个数据集共同具有的metrics"""
        # 获取两个数据集的所有列名
        mix_columns = set(self.mix_data.columns)
        single_columns = set(self.single_data.columns)
        
        # 找出共同的列名
        common_columns = mix_columns.intersection(single_columns)
        
        # 过滤掉非metrics列（如epoch, learning_rate等）
        non_metric_columns = {'epoch', 'learning_rate', 'alpha', 'train_loss', 
                             'contrastive_loss', 'feature_matching_loss', 
                             'avg_img2text_recall', 'overall_avg_recall'}
        
        # 识别recall相关的metrics
        recall_metrics = [col for col in common_columns 
                         if 'recall' in col.lower() and col not in non_metric_columns]
        
        return recall_metrics
    
    def _detect_common_min_epoch(self):
        """自动检测两个数据集共同的最小epoch"""
        mix_min_epoch = self.mix_data['epoch'].min()
        single_min_epoch = self.single_data['epoch'].min()
        return max(mix_min_epoch, single_min_epoch)  # 取较大的最小值作为共同最小值
    
    def _detect_common_max_epoch(self):
        """自动检测两个数据集共同的最大epoch"""
        mix_max_epoch = self.mix_data['epoch'].max()
        single_max_epoch = self.single_data['epoch'].max()
        return min(mix_max_epoch, single_max_epoch)  # 取较小的最大值作为共同最大值
    
    def _filter_data_by_common_epochs(self):
        """根据共同的epoch范围筛选数据"""
        self.mix_data = self.mix_data[(self.mix_data['epoch'] >= self.common_min_epoch) & 
                                     (self.mix_data['epoch'] <= self.max_epoch)]
        self.single_data = self.single_data[(self.single_data['epoch'] >= self.common_min_epoch) & 
                                           (self.single_data['epoch'] <= self.max_epoch)]
    
    def _group_metrics_by_dataset(self):
        """将metrics按数据集分组"""
        metrics_by_dataset = {}
        
        for metric in self.common_metrics:
            # 提取数据集名称（如'ImageNet_1K', 'OK-VQA'等）
            parts = metric.split('_')
            if len(parts) >= 3:
                dataset_name = '_'.join(parts[:-2])  # 去掉最后两个部分（img2text_recall@1等）
            else:
                dataset_name = 'Unknown'
            
            if dataset_name not in metrics_by_dataset:
                metrics_by_dataset[dataset_name] = []
            
            # 提取recall等级（@1, @5, @10）
            recall_level = parts[-1]
            display_name = f"{dataset_name} Recall{recall_level}"
            
            metrics_by_dataset[dataset_name].append((metric, display_name))
        
        return metrics_by_dataset
    
    def plot_training_loss_comparison(self):
        """Plot training loss comparison between mix and single methods"""
        plt.figure(figsize=(12, 6))
        
        # Plot total training loss
        plt.subplot(1, 2, 1)
        plt.plot(self.mix_data['epoch'], self.mix_data['train_loss'], 
                 'r-', marker='o', label='Mix Loss Method')
        plt.plot(self.single_data['epoch'], self.single_data['train_loss'], 
                 'b-', marker='s', label='Single Loss Method')
        plt.title('Total Training Loss Comparison')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend()
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))
        
        # Plot contrastive loss
        plt.subplot(1, 2, 2)
        plt.plot(self.mix_data['epoch'], self.mix_data['contrastive_loss'], 
                 'r-', marker='o', label='Mix Loss - Contrastive')
        plt.plot(self.single_data['epoch'], self.single_data['contrastive_loss'], 
                 'b-', marker='s', label='Single Loss - Contrastive')
        plt.title('Contrastive Loss Comparison')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend()
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))
        
        plt.tight_layout()
        plt.savefig(os.path.join(self.output_dir, 'training_loss_comparison.png'), 
                    dpi=300, bbox_inches='tight')
        plt.close()
    
    def plot_feature_matching_loss(self):
        """Plot feature matching loss comparison"""
        # 检查两个数据集是否都有feature_matching_loss列
        if 'feature_matching_loss' in self.mix_data.columns and 'feature_matching_loss' in self.single_data.columns:
            plt.figure(figsize=(10, 6))
            
            plt.plot(self.mix_data['epoch'], self.mix_data['feature_matching_loss'], 
                     'g-', marker='^', label='Mix Loss Method')
            plt.plot(self.single_data['epoch'], self.single_data['feature_matching_loss'], 
                     'y-', marker='d', label='Single Loss Method')
            plt.title('Feature Matching Loss Comparison')
            plt.xlabel('Epoch')
            plt.ylabel('Loss')
            plt.legend()
            plt.grid(True, linestyle='--', alpha=0.7)
            plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))
            
            plt.tight_layout()
            plt.savefig(os.path.join(self.output_dir, 'feature_matching_loss_comparison.png'), 
                        dpi=300, bbox_inches='tight')
            plt.close()
        else:
            print("Warning: 'feature_matching_loss' column not found in both datasets, skipping feature matching loss plot.")
    
    def plot_alpha_scheduler(self): 
        """Plot alpha scheduler for mix method"""
        if 'alpha' in self.mix_data.columns:
            plt.figure(figsize=(10, 6))
            
            plt.plot(self.mix_data['epoch'], self.mix_data['alpha'], 
                     'm-', marker='*')
            plt.title('Alpha Value Schedule (Mix Loss Method)')
            plt.xlabel('Epoch')
            plt.ylabel('Alpha Value')
            plt.grid(True, linestyle='--', alpha=0.7)
            plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))
            plt.ylim(0, 1.1)  # Set y-axis range for better visibility
            
            plt.tight_layout()
            plt.savefig(os.path.join(self.output_dir, 'alpha_scheduler.png'), 
                        dpi=300, bbox_inches='tight')
            plt.close()
        else:
            print("Warning: 'alpha' column not found in mix dataset, skipping alpha scheduler plot.")
    
    def plot_recall_metrics(self):
        """Plot recall metrics comparison using automatically detected metrics"""
        if not self.common_metrics:
            print("Warning: No common recall metrics found, skipping recall metrics plot.")
            return
        
        # 将metrics按数据集分组
        metrics_by_dataset = self._group_metrics_by_dataset()
        
        # 计算需要的子图数量
        total_plots = sum(len(metrics) for metrics in metrics_by_dataset.values())
        
        # 确定子图布局（尽可能接近方形）
        rows = int(np.ceil(np.sqrt(total_plots)))
        cols = int(np.ceil(total_plots / rows))
        
        fig, axes = plt.subplots(rows, cols, figsize=(cols*6, rows*5))
        axes = axes.flatten() if total_plots > 1 else [axes]
        
        plot_idx = 0
        for dataset_name, metrics in metrics_by_dataset.items():
            for metric_col, metric_name in metrics:
                if plot_idx < len(axes):
                    axes[plot_idx].plot(self.mix_data['epoch'], self.mix_data[metric_col], 
                                       'r-', marker='o', label='Mix Loss Method')
                    axes[plot_idx].plot(self.single_data['epoch'], self.single_data[metric_col], 
                                       'b-', marker='s', label='Single Loss Method')
                    axes[plot_idx].set_title(metric_name)
                    axes[plot_idx].set_xlabel('Epoch')
                    axes[plot_idx].set_ylabel('Recall')
                    axes[plot_idx].legend()
                    axes[plot_idx].grid(True, linestyle='--', alpha=0.7)
                    axes[plot_idx].xaxis.set_major_locator(MaxNLocator(integer=True))
                    axes[plot_idx].set_ylim(0, 1.1)  # Set y-axis range for better visibility
                    plot_idx += 1
        
        # 隐藏未使用的子图
        for i in range(plot_idx, len(axes)):
            axes[i].axis('off')
        
        plt.tight_layout()
        plt.savefig(os.path.join(self.output_dir, 'recall_metrics_comparison.png'), 
                    dpi=300, bbox_inches='tight')
        plt.close()
    
    def plot_average_metrics(self):
        """Plot average metrics comparison"""
        plt.figure(figsize=(12, 6))
        
        # Plot average image-to-text recall if available
        if 'avg_img2text_recall' in self.mix_data.columns and 'avg_img2text_recall' in self.single_data.columns:
            plt.subplot(1, 2, 1)
            plt.plot(self.mix_data['epoch'], self.mix_data['avg_img2text_recall'], 
                     'r-', marker='o', label='Mix Loss Method')
            plt.plot(self.single_data['epoch'], self.single_data['avg_img2text_recall'], 
                     'b-', marker='s', label='Single Loss Method')
            plt.title('Average Image-to-Text Recall Comparison')
            plt.xlabel('Epoch')
            plt.ylabel('Recall')
            plt.legend()
            plt.grid(True, linestyle='--', alpha=0.7)
            plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))
            plt.ylim(0.5, 0.8)  # Set y-axis range for better visibility
        else:
            plt.subplot(1, 2, 1)
            plt.text(0.5, 0.5, 'Average Image-to-Text Recall\nnot available in both datasets', 
                     horizontalalignment='center', verticalalignment='center', 
                     transform=plt.gca().transAxes)
            plt.axis('off')
        
        # Plot overall average recall if available
        if 'overall_avg_recall' in self.mix_data.columns and 'overall_avg_recall' in self.single_data.columns:
            plt.subplot(1, 2, 2)
            plt.plot(self.mix_data['epoch'], self.mix_data['overall_avg_recall'], 
                     'r-', marker='o', label='Mix Loss Method')
            plt.plot(self.single_data['epoch'], self.single_data['overall_avg_recall'], 
                     'b-', marker='s', label='Single Loss Method')
            plt.title('Overall Average Recall Comparison')
            plt.xlabel('Epoch')
            plt.ylabel('Recall')
            plt.legend()
            plt.grid(True, linestyle='--', alpha=0.7)
            plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))
            plt.ylim(0.5, 0.8)  # Set y-axis range for better visibility
        else:
            plt.subplot(1, 2, 2)
            plt.text(0.5, 0.5, 'Overall Average Recall\nnot available in both datasets', 
                     horizontalalignment='center', verticalalignment='center', 
                     transform=plt.gca().transAxes)
            plt.axis('off')
        
        plt.tight_layout()
        plt.savefig(os.path.join(self.output_dir, 'average_metrics_comparison.png'), 
                    dpi=300, bbox_inches='tight')
        plt.close()
    
    def generate_all_plots(self):
        """Generate all comparison plots"""
        print("Generating training loss comparison plot...")
        self.plot_training_loss_comparison()
        
        print("Generating feature matching loss comparison plot...")
        self.plot_feature_matching_loss()
        
        print("Generating alpha scheduler plot...")
        self.plot_alpha_scheduler()
        
        print("Generating recall metrics comparison plots...")
        self.plot_recall_metrics()
        
        print("Generating average metrics comparison plot...")
        self.plot_average_metrics()
        
        print(f"All plots have been saved to {self.output_dir}")

if __name__ == "__main__":
    # Define file paths
    mix_csv_path = 'outputs/history/new_mix/training_history.csv'
    single_csv_path = 'outputs/history/real_single/training_history.csv'
    
    # Create visualizer and generate all plots
    visualizer = ModelComparisonVisualizer(
        mix_csv_path=mix_csv_path,
        single_csv_path=single_csv_path,
        output_dir='outputs/comparison_results',
        # 不指定max_epoch，使用自动检测的值
    )
    
    # Generate all visualization plots
    visualizer.generate_all_plots()